import json
from tqdm import tqdm
from time import time
from milvus import MilvusOP

valid_set_pth = r'D:\Users\JHC258\projects\文本检索\evaluate\valid_set\t2q.json'
with open(valid_set_pth, 'r', encoding='utf-8') as f:
    text2query_ls = json.load(f)

config = {'ns': 0, 'ts': 1, 'e': 1, 'td': 1}
milvus_op = MilvusOP(db_name='state_vector_db', collection_name="hybrid2")
eval_correct_doc_docx, total_doc_docx = 0, 0
eval_correct_xls_xlsx, total_xls_xlsx = 0, 0
eval_correct_pdf, total_pdf = 0, 0
eval_correct_ofd, total_ofd = 0, 0
eval_correct_xml, total_xml = 0, 0
eval_correct_img, total_img = 0, 0
total_time = 0
correct_num = 0
block_num = len(text2query_ls)
error_res = {}
type_set = {'.doc', '.docx', '.xls', '.xlsx', '.ofd', '.xml', '.jpg', '.pdf'}
bar = tqdm(text2query_ls.items(), desc='Milvus评估')
for key, value in bar:
    query = key
    label = value['label']
    label_type = value['type']
    start_time = time()
    res = milvus_op.hybrid_search([query], limit=15, config=config)
    end_time = time()
    total_time = total_time + (end_time - start_time)
    names = {i['name'] for i in res}
    if label in names:
        correct_num = correct_num + 1
        if label_type in {'.doc', '.docx'}:
            eval_correct_doc_docx = eval_correct_doc_docx + 1
        elif label_type in {'.xls', '.xlsx'}:
            eval_correct_xls_xlsx = eval_correct_xls_xlsx + 1
        elif label_type == '.pdf':
            eval_correct_pdf = eval_correct_pdf + 1
        elif label_type == '.ofd':
            eval_correct_ofd = eval_correct_ofd + 1
        elif label_type == '.xml':
            eval_correct_xml = eval_correct_xml + 1
        else:
            eval_correct_img = eval_correct_img + 1
    else:
        error_res[key] = value
    if label_type in {'.doc', '.docx'}:
        total_doc_docx = total_doc_docx + 1
    elif label_type in {'.xls', '.xlsx'}:
        total_xls_xlsx = total_xls_xlsx + 1
    elif label_type == '.pdf':
        total_pdf = total_pdf + 1
    elif label_type == '.ofd':
        total_ofd = total_ofd + 1
    elif label_type == '.xml':
        total_xml = total_xml + 1
    else:
        total_img = total_img + 1
    bar.set_description(f'准确率：{round(correct_num / block_num, 4)}')

res_dict = {
    '准确率': correct_num / block_num if block_num != 0 else 0,
    '平均用时': total_time / block_num,
    '总数': block_num,
    '正确数': correct_num,
    '详细统计': [
        {'doc_docx准确率': eval_correct_doc_docx / total_doc_docx if total_doc_docx != 0 else 0,
         'doc_docx正确数': eval_correct_doc_docx,
         'doc_docx总数': total_doc_docx,
         },
        {'xls_xlsx准确率': eval_correct_xls_xlsx / total_xls_xlsx if total_xls_xlsx != 0 else 0,
         'xls_xlsx正确数': eval_correct_xls_xlsx,
         'xls_xlsx总数': total_xls_xlsx,
         },
        {'ofd准确率': eval_correct_ofd / total_ofd if total_ofd != 0 else 0,
         'ofd正确数': eval_correct_ofd,
         'ofd总数': total_ofd,
         },
        {'xml准确率': eval_correct_xml / total_xml if total_xml != 0 else 0,
         'xml正确数': eval_correct_xml,
         'xml总数': total_xml,
         },
        {'pdf准确率': eval_correct_pdf / total_pdf if total_pdf != 0 else 0,
         'pdf正确数': eval_correct_pdf,
         'pdf总数': total_pdf,
         },
        {'img准确率': eval_correct_img / total_img if total_img != 0 else 0,
         'img正确数': eval_correct_img,
         'img总数': total_img,
         },
    ]
}
print('结果评估: ', correct_num / len(text2query_ls))
json_path = r'D:\Users\JHC258\projects\文本检索\evaluate\eval_out\test2.json'
error_json_path = r'D:\Users\JHC258\projects\文本检索\evaluate\valid_set\t2q_error2.json'
with open(json_path, 'w+', encoding='utf-8') as f:
    json.dump(res_dict, f, ensure_ascii=False)
# with open(error_json_path, 'w+', encoding='utf-8') as f:
#     json.dump(error_res, f, ensure_ascii=False)
